import joblib
import numpy as np
from catboost import CatBoostClassifier
from pandas import read_csv
from sklearn.model_selection import train_test_split


def get_dataset(filepath):
    with open(filepath, encoding='UTF-8') as file:
        data = read_csv(file)
    return data.iloc[:, 0:16], data.iloc[:, 16:]


def train_model(train_, labels_):
    categorical_features_indices = np.where(train_.dtypes != np.float_)[0]
    model = CatBoostClassifier(cat_features=categorical_features_indices)
    model.fit(train_, labels_)
    joblib.dump(model, 'model.m')


def load_model(filepath):
    return joblib.load(filepath)


if __name__ == '__main__':
    mat, targets = get_dataset('Ob.csv')
    x_train, x_test, y_train, y_test = train_test_split(mat, targets, test_size=0.3, random_state=8)
    train_model(x_train, y_train)
    clf = load_model('model.m')
    print(clf.score(x_train, y_train))
    print(clf.score(x_test, y_test))
